from typing import List, Union
import cv2
import os.path as osp
from pathlib import Path

import cv2
import numpy as np
from PIL import Image

import insightface
from insightface.app.common import Face
from loguru import logger

class ImageCache:
    data = {}

def get_image(name, to_rgb=False):
    key = (name, to_rgb)
    if key in ImageCache.data:
        return ImageCache.data[key]
    images_dir = Path(__file__).parent.absolute()
    image_file = osp.join(images_dir, name)
    img = cv2.imread(image_file)
    if to_rgb:
        img = img[:,:,::-1]
    ImageCache.data[key] = img
    return img

def swap_face(
    source_img: str,
    target_img: str,
    output_img: str
    # source_faces_index: int = 0,
    # gender_source: int = 0,
    # gender_target: int = 0
):
    root_path = Path(__file__).parent.absolute()
    app = insightface.app.FaceAnalysis(name='buffalo_l', root=root_path)
    app.prepare(ctx_id=0, det_size=(640, 640))
    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', root=root_path, download=True, download_zip=True)

    img = get_image(target_img)
    scenes = app.get(img)

    img2 = get_image(source_img)
    face = app.get(img2)
    source_face = face[0]

    res = img.copy()
    for scene in scenes:
        res = swapper.get(img, scene, source_face, paste_back=True)
    print(output_img)
    cv2.imwrite(output_img, res)

# if __name__ == '__main__':
#     swap_face('faces/aimagine/2.pic.jpg', 'scenes/1/1-1.jpg', 'output_images/t1.jpg')